# 导入环境和学习方法
from part1.env import ArmEnv
from part1.rl import DDPG

# 设置全局变量
MAX_EPISODES = 500
MAX_EP_STEPS = 200

# 设置环境
env = ArmEnv()
s_dim = env.state_dim
a_dim = env.action_dim
a_bound = env.action_bound

# 设置rl方法
rl = DDPG(a_dim,s_dim,a_bound)

# 开始训练
for i in range(MAX_EPISODES):
    s = env.reset()
    for j in range(MAX_EP_STEPS):
        # 可视化
        env.render()

        # 选择一个动作执行
        a = rl.choose_action(s)

        # 获得执行动作后的反馈 所以把a传入获取下一个状态s_ 奖励r 是否完成done
        s_,r,done = env.step(a)

        # 存储到记忆库中去
        rl.store_tranistion(s,a,r,s_)

        # 如果存满了就开始学习
        if rl.memory_full:
            rl.learn()
        s = s_